package cn.edu.laud.sany.dataimport;

import java.io.IOException;
import java.nio.ByteBuffer;

import org.apache.cassandra.thrift.Cassandra;
import org.apache.cassandra.thrift.Column;
import org.apache.cassandra.thrift.ColumnOrSuperColumn;
import org.apache.cassandra.thrift.ConsistencyLevel;
import org.apache.cassandra.thrift.InvalidRequestException;
import org.apache.cassandra.thrift.Mutation;
import org.apache.cassandra.thrift.TimedOutException;
import org.apache.cassandra.thrift.UnavailableException;
import org.apache.cassandra.utils.ByteBufferUtil;
import org.apache.thrift.TException;
import org.apache.thrift.protocol.TBinaryProtocol;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.transport.TFramedTransport;
import org.apache.thrift.transport.TSocket;
import org.apache.thrift.transport.TTransport;
import org.apache.thrift.transport.TTransportException;

import cn.edu.laud.sany.types.RowData;
import cn.edu.laud.sany.types.UnConfiguredCFException;
import cn.edu.laud.sany.utils.MutationMap;
import cn.edu.laud.sany.utils.ProcessManager;

public class LaUDRunner_Thrift extends LaUDRunner
{
	public static int MUTATION_BATCH_SIZE = 10000;
	protected static final ConsistencyLevel CL = ConsistencyLevel.ANY;
	
	protected TTransport tf;
	protected TProtocol proto;
	protected Cassandra.Client client;
	protected MutationMap mutationMap;
	
	public LaUDRunner_Thrift(int no, String keyspace)
	{
		super(no, keyspace);
		mutationMap = new MutationMap();
	}
	
	@Override
	public boolean openConnection()
	{
		String ip = LaUDImport.IP[threadNum % LaUDImport.IP.length];
		try
		{
			logger.info("Try get Thrift connection...");
			tf = new TFramedTransport(new TSocket(ip, LaUDImport.laudPort));
			proto = new TBinaryProtocol(tf);
			client = new Cassandra.Client(proto);
			tf.open();
			client.set_keyspace(keyspace);
			logger.info("Get Thrift connection successful.");
		}
		catch (TTransportException e)
		{
			logger.error("Get Thrift connection error:" + e.toString() + ", LaUD IP:" + ip);
			return false;
		}
		catch (InvalidRequestException e)
		{
			logger.error("Get Thrift connection error:" + e.toString() + ", LaUD IP:" + ip);
			return false;
		}
		catch (TException e)
		{
			logger.error("Get Thrift connection error:" + e.toString() + ", LaUD IP:" + ip);
			return false;
		}
		return true;
	}
	
	@Override
	public void closeConnection()
	{
		try
		{
			if (tf != null)
			{
				tf.close();
			}
			logger.info("Laud connect closed successfully.");
		}
		catch (Exception e)
		{
			logger.error(e.getMessage());
			e.printStackTrace();
		}
	}
	
	public void endImport() throws IOException, UnConfiguredCFException
	{
		if (mutationMap.mutationCount > 0)
		{
			batchImport(client, mutationMap);
		}
	}
	
	@Override
	public void addRowData(RowData rowData) throws IOException, UnConfiguredCFException
	{
		Column column = new Column(ByteBufferUtil.bytes(rowData.localTime));
		column.setValue(ByteBufferUtil.bytes(rowData.paravalue));
		column.setTimestamp(System.currentTimeMillis());
		Mutation mutation = new Mutation();
		ColumnOrSuperColumn colOrSuperCol = new ColumnOrSuperColumn();
		colOrSuperCol.column = column;
		mutation.column_or_supercolumn = colOrSuperCol;
		ByteBuffer keyBuffer;
		String cfName;
		if (LaUDImport.schema == -1)
		{
			keyBuffer = ByteBufferUtil.bytes(rowData.equipCode + rowData.paraId);
			cfName = LaUDImport.cfNameFormat;
		}
		else
		{
			keyBuffer = ByteBufferUtil.bytes(rowData.equipCode);
			cfName = String.format(LaUDImport.cfNameFormat, rowData.paraId);
		}
		mutationMap.addMutation(keyBuffer, cfName, mutation);
		
		if (mutationMap.mutationCount >= MUTATION_BATCH_SIZE)
		{
			batchImport(client, mutationMap);
		}
	}
	
	protected void batchImport(Cassandra.Client nowClient, MutationMap map)
			throws UnConfiguredCFException
	{
		boolean flag = false;
		while (flag == false)
		{
			try
			{
				nowClient.batch_mutate(map.getMap(), CL);
				ProcessManager.addRealImportCount(map.mutationCount);
				flag = true;
				map.clearMutation();
			}
			catch (InvalidRequestException e)
			{
				if (e.getWhy() != null && e.getWhy().startsWith("unconfigured columnfamily"))
				{
					logger.error("One batch Import error!");
					logger.error("Caused by: " + e.toString());
					logger.error("Please check you cfs file!");
					throw (new UnConfiguredCFException());
				}
				else
				{
					logger.error("One batch import error, will try again after 1 sec!");
					logger.error("Caused by: " + e.toString());
				}
			}
			catch (UnavailableException e)
			{
				logger.error("One batch import error, will try again after 1 sec!");
				logger.error("Caused by: " + e.toString());
				WaitSecond(1);
			}
			catch (TimedOutException e)
			{
				logger.error("One batch Import time out, will try again after 1 sec !");
				WaitSecond(1);
			}
			catch (TTransportException e)
			{
				logger.error("One batch import error, will try to get connection after 5 sec!");
				logger.error("Caused by: " + e.toString());
				WaitSecond(5);
			}
			catch (TException e)
			{
				logger.error("One batch import error, will try again after 1 sec!");
				logger.error("Caused by: " + e.toString());
				WaitSecond(1);
			}
		}
		
	}
	
	private void WaitSecond(int secound)
	{
		try
		{
			Thread.sleep(secound * 1000);
		}
		catch (InterruptedException e)
		{
		}
	}
	
}
